import torch
from torch import nn
from efficientnet_pytorch import EfficientNet
from .deformable_transformer import build_deformable_transformer
import copy
import os


def _get_clones(module, N):
    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])


class DeformableTransformerMIL(nn.Module):
    """ This is the DT-MIL module that performs WSI Classification """

    def __init__(self, backbone, transformer, num_classes, c_backbone, num_queries, num_feature_levels=1):
        """ Initializes the model.
        Parameters:
            backbone: efficientNet
            num_classes: number of object classes
        """
        super().__init__()
        self.num_queries = num_queries
        self.transformer = transformer
        hidden_dim = transformer.d_model

        self.backbone = backbone
        self.input_proj = nn.Sequential(nn.Conv2d(c_backbone, hidden_dim, kernel_size=1), nn.GroupNorm(32, hidden_dim))
        self.transformer = transformer
        self.class_embed = nn.ModuleList([
            nn.MaxPool1d(num_queries),
            nn.Linear(hidden_dim, num_classes)
        ])

    def forward(self, imgs):
        features = self.backbone.extract_features(imgs)
        # bs * C * H * W
        quries = self.input_proj(features)
        # bs * dim * H * W
        patches = self.transformer(quries)
        # bs * @sum * dim
        cls_token = self.class_embed[0](patches.permute(0, 2, 1))
        # bs * dim * 1
        class_no = self.class_embed[1](cls_token.squeeze())
        return class_no


def build(args):
    num_classes = args.num_classes

    backbone = EfficientNet.from_name('efficientnet-b0')
    # print(os.listdir())
    state_dict = torch.load('data/adv-efficientnet-b0-b64d5a18.pth')
    backbone.load_state_dict(state_dict)

    transformer = build_deformable_transformer(args)
    model = DeformableTransformerMIL(
        backbone,
        transformer,
        num_classes=num_classes,
        num_queries=args.num_queries,
        num_feature_levels=args.num_feature_levels,
        c_backbone = 1280
    )

    return model
